#include "BiCGSTAB_precond_CRS.h"

void BiCGSTAB_precond_CRS(double complex *val, int *col_ind, int *row_ptr, double complex *b, double complex *x0, int n, double tolx, double tolerance)

{
  int i=0, j=0, k=0;
  int col=0;
  
  double complex lo0=0, lo1=0, beta=0, omega=0, alpha=0;
  
  //  double tolerance=1.0e-10; 
  int iter=0;
  int maxIter=5000;
  double maxm=1, maxb;
  double complex temp=0, denom=0;

  double maxm_r0=0;


  double complex *r0, *r, *x, *p, *p_hat, *v, *s, *s_hat, *t, *z;
  r0   =(double complex*)malloc(n*sizeof(double complex));
  r    =(double complex*)malloc(n*sizeof(double complex));
  x    =(double complex*)malloc(n*sizeof(double complex));
  p    =(double complex*)malloc(n*sizeof(double complex)); 
  p_hat=(double complex*)malloc(n*sizeof(double complex)); 
  v    =(double complex*)malloc(n*sizeof(double complex)); 
  s    =(double complex*)malloc(n*sizeof(double complex)); 
  s_hat=(double complex*)malloc(n*sizeof(double complex)); 
  t    =(double complex*)malloc(n*sizeof(double complex)); 
  z    =(double complex*)malloc(n*sizeof(double complex)); 

  int DILU=1, SSOR=0, ICUU=0;

  

  //ssor
  double para=1.50;
   
  //* D-ILU 
  double complex *Precon=(double complex *)malloc((row_ptr[n])*sizeof(double complex));
  int *diag_ptr=(int *)malloc(n*sizeof(int));
  double complex *pivots=(double complex *)malloc(n*sizeof(double complex));
  double complex *Kcol=(double complex *)malloc(n*sizeof(double complex));

  int found=0;
  double complex element=0;

  for(i=0;i<n;i++){
     for(j=row_ptr[i];j<row_ptr[i+1];j++){
        if(col_ind[j]==i){
          diag_ptr[i]=j; pivots[i]=val[j]; 
          break;
        }
     }
     if(cabs(pivots[i])<1.0e-16){ 
       printf("wrong in BiCGSTAB, zero diagonal\n"); 
       getchar(); return; 
     }
  }
  for(i=0;i<row_ptr[n];i++){ 
     Precon[i]=val[i];
  } 

  //D-ILU
  if(DILU==1){
    for(i=0;i<n;i++){
       pivots[i]=1.0/pivots[i];
       for(j=diag_ptr[i]+1;j<row_ptr[i+1];j++){
          found=0;
          for(k=row_ptr[col_ind[j]];k<diag_ptr[col_ind[j]];k++){
             if(col_ind[k]==i){
               found=1;
               element=Precon[k];
             }
          }
          if(found==1){
            Precon[diag_ptr[col_ind[j]]]=Precon[diag_ptr[col_ind[j]]]-element*pivots[i]*Precon[j];
          }
       }
    }
  }


  //Cholesky 
  if(ICUU==1){ 
    for(k=0;k<n;k++){
       for(j=0;j<n;j++){ 
          Kcol[j]=0;
       }
       for(j=row_ptr[k];j<row_ptr[k+1];j++){
          Kcol[col_ind[j]]=Precon[j];
       }

       for(j=diag_ptr[k]+1;j<row_ptr[k+1];j++){
          element=Precon[j];
          for(i=diag_ptr[col_ind[j]];i<row_ptr[col_ind[j]+1];i++){
             Precon[i]=Precon[i]-Kcol[col_ind[i]]*element/Precon[diag_ptr[k]];
          }
       }

       element=csqrt(Precon[diag_ptr[k]]);
       for(j=diag_ptr[k];j<row_ptr[k+1];j++){
          Precon[j]=Precon[j]/element;
       }
    }  
  }

  if(SSOR==1){
    for(i=0;i<n;i++){
       pivots[i]=1.0/pivots[i];
    }
  }

  
 

   
  // x0 initial guess
  for(k=0;k<n;k++){
     x[k]=x0[k]; 
  }

  // 1. r0=b-Ax0 
  for(i=0;i<n;i++){
     temp=0;
     for(j=row_ptr[i];j<row_ptr[i+1];j++){
        temp=temp+val[j]*x[col_ind[j]];
     }
     r[i]=b[i]-temp; 
  } 

  temp=0; 
  for(i=0;i<n;i++){
     temp=temp+conj(r[i])*r[i];
  } 
  maxm=cabs(csqrt(temp)); maxm_r0=maxm;
  printf("maxm= %E\n",maxm); 
  if(maxm<tolerance){ 
    printf("maxm= %E\n",maxm); 
    goto ending;
  }  
  if(maxm<tolerance){ 
    goto ending; 
  } 

  // 2. choose r_hat = r0  to be fixed.  
  for(j=0;j<n;j++){  
     r0[j]=r[j]; 
  }

  // 3. 
  lo0=1; alpha=1; omega=1; 

  //4.
  for(j=0;j<n;j++){  
     v[j]=0; p[j]=0; 
  }
  

  // 5. iteration
  for(iter=0;iter<maxIter;iter++){ 
     // 1. loi
     temp=0;
     for(k=0;k<n;k++){
        temp=temp+conj(r0[k])*(r[k]);
     }
     lo1=temp;  if(cabs(lo1)==0){ printf("lo =0, method fails\n"); goto ending;}

     // 2. beta // 3. p
     if(iter==0){
       for(j=0;j<n;j++){ p[j]=r[j];}
     }
     else{
       beta=(lo1/lo0)*(alpha/omega);
       for(j=0;j<n;j++){
          p[j]=r[j]+beta*(p[j]-omega*v[j]);
       }
     } 

     //3. Solve Mp_hat = p   
        //D-ILU
        if(DILU==1){
          for(i=0;i<n;i++){
             temp=0;
             for(j=row_ptr[i];j<diag_ptr[i];j++){
                temp=temp+Precon[j]*z[col_ind[j]];
             }
             z[i]=(p[i]-temp)*pivots[i];
          }
          for(i=n-1;i>=0;i--){
             temp=0;
             for(j=diag_ptr[i]+1;j<row_ptr[i+1];j++){
                temp=temp+Precon[j]*p_hat[col_ind[j]];
             } 
             p_hat[i]=z[i]-temp*pivots[i];
          }
        } 
        //SSOR
        if(SSOR==1){        
          for(i=0;i<n;i++){
             temp=0;
             for(j=row_ptr[i];j<diag_ptr[i];j++){
                temp=temp+para*Precon[j]*z[col_ind[j]];
             }
             z[i]=(para*(2-para)*p[i]-temp)*pivots[i];
          }
          for(i=n-1;i>=0;i--){
             temp=0;
             for(j=diag_ptr[i]+1;j<row_ptr[i+1];j++){
                temp=temp+para*Precon[j]*p_hat[col_ind[j]];
             } 
             //p_hat[i]=(z[i]/pivots[i]-temp)*pivots[i];
             p_hat[i]=z[i]-temp*pivots[i];
          }
        }
        //IC  
        if(ICUU==1){ 
          for(i=0;i<n;i++){ 
             Kcol[i]=p[i];
          }
          for(i=0;i<n;i++){
             z[i]=Kcol[i]/Precon[diag_ptr[i]]; 
             for(j=diag_ptr[i]+1;j<row_ptr[i+1];j++){
                Kcol[col_ind[j]]=Kcol[col_ind[j]]-z[i]*Precon[j]; 
             } 
          }
          for(i=n-1;i>=0;i--){
             temp=0;
             for(j=diag_ptr[i]+1;j<row_ptr[i+1];j++){
                temp=temp+Precon[j]*p_hat[col_ind[j]];
             } 
             p_hat[i]=(z[i]-temp)/Precon[diag_ptr[i]];
          }
        } 

     // 4. v=Ap_hat 
     for(i=0;i<n;i++){
        temp=0;
        for(j=row_ptr[i];j<row_ptr[i+1];j++){
           temp=temp+val[j]*p_hat[col_ind[j]];
        } 
        v[i]=temp; 
     }  

     // 5. alpha
     temp=0;
     for(j=0;j<n;j++){
        temp=temp+conj(r0[j])*(v[j]);
     }
     alpha=lo1/temp;  
     if(cabs(temp)==0){ printf("Error alpha_K: denom =0\n"); getchar();}   
 
     // 6. s 
     temp=0;
     for(j=0;j<n;j++){
        s[j]=r[j]-alpha*v[j];
        temp=temp+conj(s[j])*s[j];
     }
     maxm=cabs(csqrt(temp));
     printf("iter = %4d, error S= %0.10E ", iter+1, maxm); 
     //     if(maxm/maxm_r0<tolerance){
     if(maxm<tolerance){
       for(i=0;i<n;i++){ 
          x[i]=x[i]+alpha*p_hat[i];
       }
       break;
     }
     

     // 6. Solve Ms_hat=s    
        // D-ILU
        if(DILU==1){
          for(i=0;i<n;i++){
             temp=0;
             for(j=row_ptr[i];j<diag_ptr[i];j++){
                temp=temp+Precon[j]*z[col_ind[j]];
             }
             z[i]=(s[i]-temp)*pivots[i];
          }
          for(i=n-1;i>=0;i--){
             temp=0;
             for(j=diag_ptr[i]+1;j<row_ptr[i+1];j++){
                temp=temp+Precon[j]*s_hat[col_ind[j]];
             } 
             s_hat[i]=z[i]-temp*pivots[i];
          }
        } 
        //SSOR
        if(SSOR==1){        
          for(i=0;i<n;i++){
             temp=0;
             for(j=row_ptr[i];j<diag_ptr[i];j++){
                temp=temp+para*Precon[j]*z[col_ind[j]];
             }
             z[i]=(para*(2-para)*s[i]-temp)*pivots[i];
          }
          for(i=n-1;i>=0;i--){
             temp=0;
             for(j=diag_ptr[i]+1;j<row_ptr[i+1];j++){
                temp=temp+para*Precon[j]*s_hat[col_ind[j]];
             } 
             //s_hat[i]=(z[i]/pivots[i]-temp)*pivots[i];
             s_hat[i]=z[i]-temp*pivots[i];
          }
        }       
        //IC  
        if(ICUU==1){  
          for(i=0;i<n;i++){ 
             Kcol[i]=s[i];
          }
          for(i=0;i<n;i++){
             z[i]=Kcol[i]/Precon[diag_ptr[i]]; 
             for(j=diag_ptr[i]+1;j<row_ptr[i+1];j++){
                Kcol[col_ind[j]]=Kcol[col_ind[j]]-z[i]*Precon[j]; 
             } 
          }
          for(i=n-1;i>=0;i--){
             temp=0;
             for(j=diag_ptr[i]+1;j<row_ptr[i+1];j++){
                temp=temp+Precon[j]*s_hat[col_ind[j]];
             } 
             s_hat[i]=(z[i]-temp)/Precon[diag_ptr[i]];
          }
        }

     // 7. t=As_hat
     for(i=0;i<n;i++){
        temp=0;
        for(j=row_ptr[i];j<row_ptr[i+1];j++){
           temp=temp+val[j]*s_hat[col_ind[j]];
        } 
        t[i]=temp;
     }  
 
     // 8. omega
     temp=0; denom=0;
     for(j=0;j<n;j++){ 
        temp=temp+conj(t[j])*(s[j]);
        denom=denom+t[j]*conj(t[j]);
     }
     omega=temp/denom; if(cabs(denom)==0){ printf("Error alpha_K: denom =0\n"); getchar();}    


     // 9. x
     maxm = 0;
     for(i=0;i<n;i++){ 
        x[i] = x[i] + alpha*p_hat[i] + omega*s_hat[i];
	maxb = cabs(alpha*p_hat[i]+omega*s_hat[i]);
        maxm = maxb > maxm ? maxb : maxm; 
     } 
     printf("iter = %4d, error X= %0.10E, ", iter+1, maxm); 
     //if(maxm/maxm_r0<tolerance){  
     //  break;
     //}
     
     // 11. r 
     temp=0;
     for(j=0;j<n;j++){ 
        r[j]=s[j]-omega*t[j];
        temp=temp+conj(r[j])*r[j]; 
     } 
     maxm=cabs(csqrt(temp));

     printf("r = %0.10E\n",maxm); 
     //     if(maxm/maxm_r0<tolerance){  
     if(maxm < tolerance){  
       break;
     }      
 
     lo0=lo1; 
     
  } 
   

ending:

  for(k=0;k<n;k++){
     x0[k]=x[k];
  }

  free(r0); 
  free(r); 
  free(x); 
  free(p); 
  free(p_hat); 
  free(v); 
  free(s); 
  free(s_hat); 
  free(t); 
  free(z);

  free(Precon); 
  free(diag_ptr); 
  free(pivots);  
  free(Kcol);

 


  return;
} 
       
